import torch
from torch.utils.data import Dataset
from get_dict import get_data
from until import EOS_token

# 调用获取数据函数
english_word2index, english_index2word, english_word_n, \
    french_word2index, french_index2word, french_word_n, my_pairs = get_data()


# 构建dataset
class MyPairDataset(Dataset):
    def __init__(self, my_pairs):
        super(MyPairDataset, self).__init__()
        self.my_pairs = my_pairs
        self.sample_len = len(my_pairs)

    def __len__(self):
        return self.sample_len

    def __getitem__(self, index):
        index = min(max(0, index), self.sample_len - 1)

        # 获取英文输入和法文输出
        english_x = my_pairs[index][0]
        french_y = my_pairs[index][1]

        # 构建输入序列
        x = [english_word2index[word] for word in english_x.split(' ')]
        # 添加结束符
        x.append(EOS_token)
        # 转换为张量
        torch_x = torch.tensor(x, dtype=torch.long)

        # 构建输出序列
        y = [french_word2index[word] for word in french_y.split(' ')]
        # 添加结束符
        y.append(EOS_token)
        # 转换为张量
        torch_y = torch.tensor(y, dtype=torch.long)

        return torch_x, torch_y


if __name__ == '__main__':
    # 创建dataset并打印一条数据
    my_dataset = MyPairDataset(my_pairs)
    print(my_dataset[0])

    # 查看此时的x的形状
    print(my_dataset[0][0].shape)
